using Distributed
addprocs(3)

@everywhere using Interpolations, Parameters, Distributions, DataFrames, CSV, LinearAlgebra, Statistics, Random, Profile, NLopt, SharedArrays
@everywhere include("mf_setup.jl")
@everywhere include("mf_valfunc.jl")
@everywhere include("mf_estimation.jl")
@everywhere include("mf_simulation.jl")
@everywhere include("mf_readin.jl")

#change the subsequent line as appropriate 
@everywhere dir = "C:\\Users\\Garrett\\Documents\\Work\\Papers\\Mig_Fam\\Model"
@everywhere cd(dir)

########################
####estimation routine

#=
This code will re-estimate the model and compute standard errors. Note that, unless you use more cores, this will likely take the better part of a week



=#
########################
using Optim, NLSolversBase

####estimate for each race
for r = 1:3
    #set up wrapper function
    dir = "C:\\Users\\Garrett\\Documents\\Work\\Papers\\Mig_Fam\\Model" ##change as appropriate 
    
    
    function wrapper(guess::Array{Float64}, grad::Vector)
        err = Objective_function(guess, dir; race = r)
        err
    end    
    opt = Opt(:LN_SBPLX, 49)
    opt.xtol_rel = 1e-4
    opt.ftol_rel = 1e-4
    opt.min_objective = wrapper
    
    #read in estimate guess
    tempthing = CSV.read("$dir\\estimates_$r.csv", DataFrame; skipto=1)
    tempthing2 = Array{Float64,2}(tempthing)'
    tempthing3 = zeros(length(tempthing2))
    for iter = 1:length(tempthing2)
        tempthing3[iter] = tempthing2[iter]
    end
    guess_init = copy(tempthing3)
    minx = guess_init 
    
    #find minimum
    (minf, minx, ret) = NLopt.optimize(opt, guess_init)
    CSV.write("$dir\\estimates_$r.csv", DataFrame(minx', :auto), header=false) #write CSV file
    
    ###Obtain standard errors
    func = TwiceDifferentiable(vars->Objective_function(vars, dir; race = r), minx) #differentiable version of OF
    hess = hessian!(func, minx)
    CSV.write("$dir\\hess_$r.csv", DataFrame(hess, :auto), header=false) #write CSV file

    #In estimation, an error was encountered where the wage effect for the New England divison for Blacks was unidentified due to
    #no Blacks in the New England ever being observed. As a workaround at the time, we used estimates for the overall sample and for Whites to form 
    #a reasonable guess of this wage parameter and its standard error. This workaround was missed when preparing the final submission.

    #A better approach would have been to fix the Black New England wage estimate at a reasonable value and re-estimate the remaining 48 paraemters 
    #from the start. When doing this, the correlation between the parameter values used in the paper and the re-estimated values is 0.998, with comparably non-existent
    #changes to the main simulation results. Nonetheless, we apologize for this error.

    var_cov_mat = pinv(hess)
    stand_errs = diag(var_cov_mat)
    CSV.write("$dir\\ses_$r.csv", DataFrame(stand_errs', :auto), header=false) #write CSV file 
end



#####No G version
function wrapper(guess::Array{Float64}, grad::Vector)
    
    guess_mod = zeros(49)
    guess_mod[1:10] = guess[1:10]
    guess_mod[11] = 0.0
    guess_mod[12] = 0.0
    guess_mod[13:24] = guess[11:22]
    guess_mod[25] = 0.0
    guess_mod[26:49] = guess[23:46]

    err = Objective_function(guess_mod, dir; race = 1)
    err
end

opt = Opt(:LN_SBPLX, 46)
opt.xtol_rel = 1e-4
opt.ftol_rel = 1e-4
opt.min_objective = wrapper

tempthing = CSV.read("$dir\\estimates_no_g.csv", DataFrame; skipto=1)
tempthing2 = Array{Float64,2}(tempthing)'
tempthing3 = zeros(length(tempthing2))
for iter = 1:length(tempthing2)
    tempthing3[iter] = tempthing2[iter]
end
guess_init = copy(tempthing3)
println(round.(guess_init, digits=4))

(minf, minx, ret) = optimize(opt, guess_init)
println("guess_init = ", round.(minx, digits = 5))
println(minf)
CSV.write("$dir\\estimates_no_g.csv", DataFrame(minx', :auto), header=false) #write CSV file










####end of file